#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Verifier.h>
#include <llvm/Support/raw_ostream.h>

using namespace llvm;

static Function *createAdd(Module &module, LLVMContext &ctx) {
  IRBuilder<> builder(ctx);
  auto *int32Ty = builder.getInt32Ty();
  auto *fnType = FunctionType::get(int32Ty, {int32Ty, int32Ty}, false);
  auto *fn = Function::Create(fnType, Function::InternalLinkage, "add", module);

  auto argsItr = fn->arg_begin();
  Argument *a = argsItr++;
  a->setName("a");
  Argument *b = argsItr++;
  b->setName("b");

  BasicBlock *entry = BasicBlock::Create(ctx, "entry", fn);
  builder.SetInsertPoint(entry);
  Value *sum = builder.CreateNSWAdd(a, b, "sum");
  builder.CreateRet(sum);
  return fn;
}

static Function *createClamp(Module &module, LLVMContext &ctx) {
  IRBuilder<> builder(ctx);
  auto *int32Ty = builder.getInt32Ty();
  auto *fnType = FunctionType::get(int32Ty, {int32Ty}, false);
  auto *fn =
      Function::Create(fnType, Function::InternalLinkage, "clamp_non_negative",
                       module);

  Argument &x = *fn->arg_begin();
  x.setName("x");

  BasicBlock *entry = BasicBlock::Create(ctx, "entry", fn);
  builder.SetInsertPoint(entry);
  Value *isNegative = builder.CreateICmpSLT(&x, builder.getInt32(0), "is.neg");
  Value *sel =
      builder.CreateSelect(isNegative, builder.getInt32(0), &x, "clamped");
  builder.CreateRet(sel);
  return fn;
}

static Function *createCompute(Module &module, LLVMContext &ctx, Function *addFn,
                               Function *clampFn) {
  IRBuilder<> builder(ctx);
  auto *int32Ty = builder.getInt32Ty();
  auto *fnType =
      FunctionType::get(int32Ty, {builder.getInt1Ty(), int32Ty}, false);
  auto *fn =
      Function::Create(fnType, Function::InternalLinkage, "compute", module);

  auto argsItr = fn->arg_begin();
  Argument *boost = argsItr++;
  boost->setName("boost");
  Argument *value = argsItr++;
  value->setName("value");

  BasicBlock *entry = BasicBlock::Create(ctx, "entry", fn);
  builder.SetInsertPoint(entry);
  Value *base = builder.CreateCall(clampFn, {value}, "base");
  Value *doubleBase =
      builder.CreateShl(base, builder.getInt32(1), "double.base");
  Value *sel =
      builder.CreateSelect(boost, doubleBase, base, "maybe.boosted");
  builder.CreateRet(sel);
  return fn;
}

static Function *createMain(Module &module, LLVMContext &ctx, Function *addFn,
                            Function *computeFn) {
  IRBuilder<> builder(ctx);
  auto *int32Ty = builder.getInt32Ty();
  auto *fnType = FunctionType::get(int32Ty, false);
  auto *fn =
      Function::Create(fnType, Function::ExternalLinkage, "main", module);

  BasicBlock *entry = BasicBlock::Create(ctx, "entry", fn);
  builder.SetInsertPoint(entry);

  Value *callAdd =
      builder.CreateCall(addFn, {builder.getInt32(4), builder.getInt32(-1)},
                         "add.result");
  Value *callComputeBoosted =
      builder.CreateCall(computeFn, {builder.getTrue(), builder.getInt32(-3)},
                         "compute.boosted");
  Value *callComputePlain =
      builder.CreateCall(computeFn,
                         {builder.getFalse(), builder.getInt32(5)},
                         "compute.plain");

  Value *sum1 =
      builder.CreateNSWAdd(callAdd, callComputeBoosted, "sum.partial");
  Value *sum2 =
      builder.CreateNSWAdd(sum1, callComputePlain, "sum.final");
  builder.CreateRet(sum2);
  return fn;
}

int main() {
  LLVMContext ctx;
  Module module("irbuilder-stage03", ctx);

  Function *addFn = createAdd(module, ctx);
  Function *clampFn = createClamp(module, ctx);
  Function *computeFn = createCompute(module, ctx, addFn, clampFn);
  createMain(module, ctx, addFn, computeFn);

  if (verifyModule(module, &errs())) {
    errs() << "Module verification failed\n";
    return 1;
  }

  module.print(outs(), nullptr);
  return 0;
}
